package com.lagou;

import javax.sql.DataSource;
import java.sql.Connection;
import java.util.HashMap;
import java.util.Map;

/**
 * @author: wudong
 * @create: 2021-05-31 14:32
 **/
public class DataSourceUtil {
    private static ThreadLocal<Map<DataSource, ConnectionHolder>> dataSourceThread = new ThreadLocal<>();

    public static void bindConnect(DataSource dataSource,ConnectionHolder connection){
        Map<DataSource, ConnectionHolder> dataSourceConnectionMap = dataSourceThread.get();
        if(dataSourceConnectionMap == null){
            dataSourceConnectionMap = new HashMap<>();
            dataSourceThread.set(dataSourceConnectionMap);
        }
        dataSourceConnectionMap.put(dataSource,connection);
    }

    public static ConnectionHolder getConnect(DataSource dataSource){
        Map<DataSource, ConnectionHolder> dataSourceConnectionMap = dataSourceThread.get();
        if(dataSourceConnectionMap == null){
            return null;
        }
        return dataSourceConnectionMap.get(dataSource);
    }

    public static void unBindConnect(DataSource dataSource){
        Map<DataSource, ConnectionHolder> dataSourceConnectionMap = dataSourceThread.get();
        if(dataSourceConnectionMap == null){
            return;
        }
        dataSourceConnectionMap.remove(dataSource);
    }

    public static interface ConnectionHolder{
        public Connection getConect();
    }
}
